import os
import random
from cognlp.utils.util import load_json


class TrexReLoader:
    def __init__(self):
        self.headers = ['sentence', 'label']
        self.label_set = set()

    def _load(self, path):
        dataset = load_json(path)
        for data in dataset:
            triples = data['triples']
            for triple in triples:
                self.label_set.add(triple['predicate']['uri'])
        return dataset

    def load(self, path):
        train_set = self._load(os.path.join(path, 'train.txt'))
        dev_set = self._load(os.path.join(path, 'dev.txt'))
        test_set = self._load(os.path.join(path, 'test.txt'))
        return train_set, dev_set, test_set

    def load_all(self, path):
        datasets = []
        for f in os.listdir(path):
            if random.random() > 0.5:
                dataset = load_json(os.path.join(path, f))
                for data in dataset:
                    triples = data['triples']
                    for triple in triples:
                        self.label_set.add(triple['predicate']['uri'])
                datasets.extend(dataset)
        return datasets

    def load_train(self, path):
        train_set = self._load(os.path.join(path, 'train.json'))
        return train_set

    def load_data(self, file):
        data_set = self._load(file)
        return data_set

    def get_labels(self):
        labels = list(self.label_set)
        labels.sort()
        return labels


class FewRelLoader:
    def __init__(self):
        self.label_set = set()

    def _load(self, path):
        data = load_json(path)
        dataset = []
        for key, value in data.items():
            self.label_set.add(key)
            for example in value:
                item = {'tokens': example['tokens'], 'h': example['h'], 't': example['t'], 'label': key}
                dataset.append(item)
        return dataset

    def load(self, path):
        train_set = self._load(os.path.join(path, 'train_wiki.json'))
        dev_set = self._load(os.path.join(path, 'val_wiki.json'))
        return train_set, dev_set

    def get_labels(self):
        labels = list(self.label_set)
        labels.sort()
        return labels
